

import mitmproxy.http
from mitmproxy import ctx
import json
import os
from datetime import datetime
import base64


class HttpTrafficExporter:
    def __init__(self):
        # 输出文件路径（可根据需要修改）
        self.output_file = self.get_output_path()
        self.request_count = 0
        self.init_output_file()

    def get_output_path(self):
        """获取输出文件路径，兼容Windows和Linux"""
        if os.name == 'nt':  # Windows
            # 优先使用桌面，如果不存在则使用当前目录
            desktop = os.path.join(os.path.expanduser("~"), "Desktop")
            if os.path.exists(desktop):
                return os.path.join(desktop, "mitmproxy_traffic_export.txt")
            else:
                return "mitmproxy_traffic_export.txt"
        else:  # Linux/Mac
            return os.path.expanduser("~/mitmproxy_traffic_export.txt")

    def init_output_file(self):
        """初始化输出文件"""
        try:
            with open(self.output_file, "w", encoding="utf-8") as f:
                f.write("=" * 80 + "\n")
                f.write("mitmproxy HTTP流量完整记录\n")
                f.write(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"输出路径: {os.path.abspath(self.output_file)}\n")
                f.write("=" * 80 + "\n\n")
            ctx.log.info(f"HTTP流量将保存到: {os.path.abspath(self.output_file)}")
        except Exception as e:
            ctx.log.error(f"无法创建输出文件: {e}")

    def request(self, flow: mitmproxy.http.HTTPFlow) -> None:
        """处理HTTP请求（立即记录）"""
        try:
            self.request_count += 1
            self.save_request_flow(flow)
            ctx.log.info(f"已记录第 {self.request_count} 个请求: {flow.request.method} {flow.request.pretty_url}")
        except Exception as e:
            ctx.log.error(f"记录请求失败: {e}")

    def response(self, flow: mitmproxy.http.HTTPFlow) -> None:
        """处理HTTP响应（更新记录）"""
        try:
            self.update_response_flow(flow)
            ctx.log.info(
                f"已更新响应: {flow.request.method} {flow.request.pretty_url} -> {flow.response.status_code if flow.response else 'No Response'}")
        except Exception as e:
            ctx.log.error(f"更新响应失败: {e}")

    def save_request_flow(self, flow: mitmproxy.http.HTTPFlow):
        """立即保存请求信息"""
        try:
            with open(self.output_file, "a", encoding="utf-8") as f:
                f.write("=" * 80 + "\n")
                f.write(f"交互 #{self.request_count}\n")
                f.write(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"URL: {flow.request.pretty_url}\n")
                f.write("=" * 80 + "\n\n")

                # 1. 原生请求信息
                self.write_request_info(f, flow)

                # 2. curl命令
                self.write_curl_command(f, flow)

                # 3. 占位符，等待响应
                f.write("📥 原生响应信息\n")
                f.write("-" * 40 + "\n")
                f.write("⏳ 等待响应中...\n\n")

                f.write("\n" + "=" * 80 + "\n\n")
        except Exception as e:
            ctx.log.error(f"保存请求信息失败: {e}")

    def update_response_flow(self, flow: mitmproxy.http.HTTPFlow):
        """更新响应信息（在原文件中找到对应请求并更新）"""
        # 由于文件追加的特性，这里我们简单地追加响应信息
        # 在实际使用中，可以通过请求ID来关联
        try:
            with open(self.output_file, "a", encoding="utf-8") as f:
                f.write(f"📥 响应更新 - {flow.request.method} {flow.request.pretty_url}\n")
                f.write("-" * 60 + "\n")
                if flow.response:
                    f.write(f"状态码: {flow.response.status_code} {flow.response.reason}\n")
                    f.write(f"响应时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")

                    # 响应头
                    f.write("\n📋 响应头:\n")
                    for name, value in flow.response.headers.items():
                        f.write(f"  {name}: {value}\n")

                    # 响应体（简化版）
                    f.write(f"\n📦 响应体长度: {len(flow.response.content) if flow.response.content else 0} bytes\n")
                    if flow.response.content:
                        content_type = flow.response.headers.get("content-type", "").lower()
                        if "application/json" in content_type and len(flow.response.content) < 1000:
                            try:
                                json_data = json.loads(flow.response.content.decode('utf-8'))
                                json_str = json.dumps(json_data, ensure_ascii=False, indent=2)
                                f.write(f"响应内容:\n{json_str}\n")
                            except:
                                f.write("响应内容: [JSON解析失败]\n")
                        else:
                            f.write("响应内容: [过大或非文本内容，已跳过]\n")
                else:
                    f.write("❌ 无响应\n")
                f.write("\n" + "=" * 60 + "\n\n")
        except Exception as e:
            ctx.log.error(f"更新响应信息失败: {e}")

    def save_complete_flow(self, flow: mitmproxy.http.HTTPFlow):
        """保存完整的请求和响应信息（兼容方法）"""
        try:
            with open(self.output_file, "a", encoding="utf-8") as f:
                f.write("=" * 80 + "\n")
                f.write(f"交互 #{self.request_count}\n")
                f.write(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"URL: {flow.request.pretty_url}\n")
                f.write("=" * 80 + "\n\n")

                # 1. 原生请求信息
                self.write_request_info(f, flow)

                # 2. 原生响应信息
                self.write_response_info(f, flow)

                # 3. curl命令
                self.write_curl_command(f, flow)

                f.write("\n" + "=" * 80 + "\n\n")
        except Exception as e:
            ctx.log.error(f"保存流量信息失败: {e}")

    def write_request_info(self, f, flow: mitmproxy.http.HTTPFlow):
        """写入原生请求信息"""
        f.write("📤 原生请求信息\n")
        f.write("-" * 40 + "\n")

        # 请求行
        f.write(f"请求方法: {flow.request.method}\n")
        f.write(f"请求URL: {flow.request.pretty_url}\n")
        f.write(f"HTTP版本: {flow.request.http_version}\n")
        f.write(f"请求路径: {flow.request.path}\n")

        # 请求头
        f.write("\n📋 请求头:\n")
        for name, value in flow.request.headers.items():
            f.write(f"  {name}: {value}\n")

        # 请求体
        f.write("\n📦 请求体:\n")
        if flow.request.content:
            content_type = flow.request.headers.get("content-type", "").lower()
            try:
                if "application/json" in content_type:
                    # JSON格式化
                    json_data = json.loads(flow.request.content.decode('utf-8'))
                    json_str = json.dumps(json_data, ensure_ascii=False, indent=2)
                    f.write(f"  内容类型: JSON\n")
                    f.write(f"  内容长度: {len(flow.request.content)} bytes\n")
                    f.write(f"  内容:\n{json_str}\n")
                elif "application/x-www-form-urlencoded" in content_type:
                    # 表单数据
                    form_data = flow.request.content.decode('utf-8')
                    f.write(f"  内容类型: 表单数据\n")
                    f.write(f"  内容长度: {len(flow.request.content)} bytes\n")
                    f.write(f"  内容: {form_data}\n")
                elif "text/" in content_type or "application/xml" in content_type:
                    # 文本数据
                    text_data = flow.request.content.decode('utf-8', errors='ignore')
                    f.write(f"  内容类型: 文本\n")
                    f.write(f"  内容长度: {len(flow.request.content)} bytes\n")
                    if len(text_data) > 2000:
                        f.write(f"  内容: {text_data[:2000]}...[截断]\n")
                    else:
                        f.write(f"  内容: {text_data}\n")
                else:
                    # 二进制数据
                    f.write(f"  内容类型: 二进制数据\n")
                    f.write(f"  内容长度: {len(flow.request.content)} bytes\n")
                    if len(flow.request.content) <= 100:
                        b64_data = base64.b64encode(flow.request.content).decode('ascii')
                        f.write(f"  内容(Base64): {b64_data}\n")
                    else:
                        f.write(f"  内容: [二进制数据过大，已跳过]\n")
            except Exception as e:
                f.write(f"  内容解析失败: {str(e)}\n")
                f.write(f"  原始长度: {len(flow.request.content)} bytes\n")
        else:
            f.write("  无请求体\n")

        f.write("\n")

    def write_response_info(self, f, flow: mitmproxy.http.HTTPFlow):
        """写入原生响应信息"""
        f.write("📥 原生响应信息\n")
        f.write("-" * 40 + "\n")

        if flow.response:
            # 响应状态
            f.write(f"状态码: {flow.response.status_code}\n")
            f.write(f"状态描述: {flow.response.reason}\n")
            f.write(f"HTTP版本: {flow.response.http_version}\n")

            # 响应头
            f.write("\n📋 响应头:\n")
            for name, value in flow.response.headers.items():
                f.write(f"  {name}: {value}\n")

            # 响应体
            f.write("\n📦 响应体:\n")
            if flow.response.content:
                content_type = flow.response.headers.get("content-type", "").lower()
                try:
                    if "application/json" in content_type:
                        # JSON响应
                        json_data = json.loads(flow.response.content.decode('utf-8'))
                        json_str = json.dumps(json_data, ensure_ascii=False, indent=2)
                        f.write(f"  内容类型: JSON\n")
                        f.write(f"  内容长度: {len(flow.response.content)} bytes\n")
                        if len(json_str) > 3000:
                            f.write(f"  内容: {json_str[:3000]}...[截断]\n")
                        else:
                            f.write(f"  内容:\n{json_str}\n")
                    elif "text/" in content_type or "application/xml" in content_type:
                        # 文本响应
                        text_data = flow.response.content.decode('utf-8', errors='ignore')
                        f.write(f"  内容类型: 文本\n")
                        f.write(f"  内容长度: {len(flow.response.content)} bytes\n")
                        if len(text_data) > 3000:
                            f.write(f"  内容: {text_data[:3000]}...[截断]\n")
                        else:
                            f.write(f"  内容: {text_data}\n")
                    else:
                        # 二进制响应
                        f.write(f"  内容类型: 二进制数据\n")
                        f.write(f"  内容长度: {len(flow.response.content)} bytes\n")
                        if len(flow.response.content) <= 200:
                            b64_data = base64.b64encode(flow.response.content).decode('ascii')
                            f.write(f"  内容(Base64): {b64_data}\n")
                        else:
                            f.write(f"  内容: [二进制数据过大，已跳过]\n")
                except Exception as e:
                    f.write(f"  内容解析失败: {str(e)}\n")
                    f.write(f"  原始长度: {len(flow.response.content)} bytes\n")
            else:
                f.write("  无响应体\n")
        else:
            f.write("❌ 未收到响应或响应为空\n")

        f.write("\n")

    def write_curl_command(self, f, flow: mitmproxy.http.HTTPFlow):
        """写入等价的curl命令"""
        f.write("🔄 等价curl命令\n")
        f.write("-" * 40 + "\n")

        curl_command = self.flow_to_curl(flow)
        f.write(curl_command + "\n")

    def flow_to_curl(self, flow: mitmproxy.http.HTTPFlow) -> str:
        """将HTTP请求转换为curl命令"""
        method = flow.request.method
        url = flow.request.pretty_url

        # 构建curl命令
        curl_parts = ["curl"]

        # 添加请求方法
        if method != "GET":
            curl_parts.extend(["-X", method])

        # 添加请求头
        for name, value in flow.request.headers.items():
            # 跳过一些自动生成的头部
            if name.lower() in ['host', 'content-length']:
                continue
            curl_parts.extend(["-H", f"'{name}: {value}'"])

        # 添加请求体
        if flow.request.content:
            content_type = flow.request.headers.get("content-type", "").lower()
            try:
                if "application/json" in content_type:
                    # JSON数据
                    json_data = json.loads(flow.request.content.decode('utf-8'))
                    json_str = json.dumps(json_data, ensure_ascii=False, separators=(',', ':'))
                    curl_parts.extend(["-d", f"'{json_str}'"])
                elif "application/x-www-form-urlencoded" in content_type:
                    # 表单数据
                    form_data = flow.request.content.decode('utf-8')
                    curl_parts.extend(["-d", f"'{form_data}'"])
                else:
                    # 其他类型的数据
                    data = flow.request.content.decode('utf-8', errors='ignore')
                    if len(data) > 1000:
                        data = data[:1000] + "...[truncated]"
                    curl_parts.extend(["-d", f"'{data}'"])
            except Exception:
                curl_parts.append("# [请求体包含二进制数据，请查看上方原生请求信息]")

        # 添加其他有用选项
        curl_parts.extend([
            "-k",  # 忽略SSL证书错误
            "--compressed",  # 支持压缩
            "--location",  # 跟随重定向
            "-v"  # 详细输出（显示头信息）
        ])

        # 最后添加URL
        curl_parts.append(f"'{url}'")

        return " \\\n  ".join(curl_parts)

    def done(self):
        """脚本结束时调用"""
        ctx.log.info(f"导出完成！共记录了 {self.request_count} 个完整HTTP交互")
        ctx.log.info(f"完整流量记录已保存到: {os.path.abspath(self.output_file)}")


# 创建导出器实例
addons = [HttpTrafficExporter()]